package org.kiwi.data.config;

import org.kiwi.data.RoutingDataSource;
import org.kiwi.data.preheat.DatabasePreheatProvider;
import org.kiwi.data.util.BeanDefinitionParserUtil;
import org.kiwi.data.vo.DataSourceMetaVO;
import org.springframework.beans.factory.config.BeanDefinition;
import org.springframework.beans.factory.config.RuntimeBeanReference;
import org.springframework.beans.factory.support.ManagedMap;
import org.springframework.beans.factory.support.RootBeanDefinition;
import org.springframework.beans.factory.xml.BeanDefinitionParser;
import org.springframework.beans.factory.xml.ParserContext;
import org.w3c.dom.Element;

import java.util.LinkedHashMap;
import java.util.Map;

import static org.kiwi.data.cons.DataConstant.*;

/**
 * Created by jack on 16/10/19.
 */
public class DataSourceBeanDefinitionParser implements BeanDefinitionParser {

    private Map<String, ManagedMap<String, RuntimeBeanReference>> groupDsMap = new LinkedHashMap<>();
    private Map<String, String> writeDsBeanNameMap = new LinkedHashMap<>();

    public BeanDefinition parse(Element element, ParserContext parserContext) {
        try {
            DataSourceMetaVO dataSourceMetaVO = BeanDefinitionParserUtil.resolve(element, parserContext);

            validateDataSourceMetaVO(dataSourceMetaVO);

            registerConfigDataSource(parserContext, dataSourceMetaVO);

            registerRoutingDataSource(parserContext, dataSourceMetaVO);

            registerPreheatProvider(parserContext);

            return null;
        } catch (Exception e) {
            parserContext.getReaderContext().error(
                    "register DataSourceBeanDefinition has some problem:[" + e.getMessage() + "]",
                    parserContext.extractSource(element));
            throw e;
        }
    }

    private void validateDataSourceMetaVO(DataSourceMetaVO dataSourceMetaVO) {
        BeanDefinitionParserUtil.validate(dataSourceMetaVO);
    }

    private void registerConfigDataSource(ParserContext parserContext, DataSourceMetaVO dataSourceMetaVO) {
        BeanDefinitionParserUtil.registerConfigDataSource(
                parserContext,
                dataSourceMetaVO,
                this.groupDsMap,
                this.writeDsBeanNameMap
        );
    }

    private void registerRoutingDataSource(ParserContext parserContext, DataSourceMetaVO dataSourceMetaVO) {
        RootBeanDefinition routingDataSource = new RootBeanDefinition(RoutingDataSource.class);
        ManagedMap<String, RuntimeBeanReference> targetDataSources = new ManagedMap<>();
        for (ManagedMap<String, RuntimeBeanReference> dsMap : this.groupDsMap.values()) {
            targetDataSources.putAll(dsMap);
        }
        routingDataSource.getPropertyValues().add(TARGET_DATA_SOURCES, targetDataSources);
        routingDataSource.getPropertyValues().add(DEFAULT_TARGET_DATA_SOURCE, new RuntimeBeanReference(this.writeDsBeanNameMap.values().iterator().next()));
        parserContext.getRegistry().registerBeanDefinition(dataSourceMetaVO.getId(), routingDataSource);
    }

    private void registerPreheatProvider(ParserContext parserContext) {
        RootBeanDefinition databasePreheatProvider = new RootBeanDefinition(DatabasePreheatProvider.class);
        parserContext.getRegistry().registerBeanDefinition(DatabasePreheatProvider.class.getSimpleName(), databasePreheatProvider);
    }

}
